#!/bin/bash
set -euo pipefail

# 脚本名称：vmount.sh

# 功能说明：
# 本脚本用于挂载虚拟磁盘或物理磁盘，并进入系统环境。

# 主要流程包括：
# 1. 安装必要的依赖工具（如 qemu-img、qemu-nbd、parted 等）。
# 2. 根据用户输入判断设备是物理磁盘还是虚拟磁盘镜像。
# 3. 如果是虚拟磁盘镜像，则使用 qemu-nbd 挂载到一个空闲的 nbd 设备上。
# 4. 挂载所有分区到指定目录。
# 5. 如果找到根文件系统，则切换到 chroot 环境。
# 6. 如果未找到根文件系统，则进入 bash 交互环境。
# 7. 使用 trap 在退出时自动断开 nbd 挂载（如果有的话）。

# 注意事项：
# - 需要 root 权限运行部分命令，如 qemu-nbd、parted。
# - 在执行磁盘操作前，建议备份重要数据，以防万一。

# 全局变量
MOUNT_POINTS=()  # 存储所有已挂载的分区
NBD_DEVICE=""
IS_VIRTUAL=false  # 标记是否为虚拟磁盘
DISK=""
TARGET_DEVICE=""
VIRTUAL_DISK_ATTACHED=false
SCRIPT_DIR=$(dirname "$(realpath "$0")")
VDISK_MOUNT_POINT="$SCRIPT_DIR/vdisk"

# 安装依赖包（仅限Ubuntu）
install_dependencies() {
    local required=("qemu-img" "qemu-nbd" "partprobe")
    local missing=()

    # 检测缺失的命令
    for cmd in "${required[@]}"; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            missing+=("$cmd")
        fi
    done

    if [ ${#missing[@]} -eq 0 ]; then
        return 0
    fi

    echo "检测到缺少依赖命令: ${missing[*]}，尝试自动安装..."
    sudo apt-get update -qq
    echo "正在安装 qemu-utils parted..."
    sudo apt-get install -y -qq qemu-utils parted >/dev/null
}

# 清理函数：如果虚拟磁盘已挂载，则断开 nbd 设备
cleanup() {
    echo "清理中..."
    # 卸载所有挂载点
    for mount_point in "${MOUNT_POINTS[@]}"; do
        echo "卸载 $mount_point..."
        sudo umount "$mount_point" 2>/dev/null
        sudo rmdir "$mount_point" 2>/dev/null
    done

    # 断开NBD连接（仅虚拟磁盘）
    if $VIRTUAL_DISK_ATTACHED && [ -n "$NBD_DEVICE" ]; then
        echo "断开 qemu-nbd..."
        sudo qemu-nbd -d "$NBD_DEVICE" >/dev/null 2>&1
    fi
}
trap cleanup EXIT SIGINT SIGTERM

# 查找空闲的 nbd 设备
select_free_nbd() {
    for dev in /dev/nbd*; do
        if [ -z "$(lsblk -n "$dev" | awk 'NR>1')" ]; then
            echo "$dev"
            return 0
        fi
    done
    return 1
}

# 检查磁盘映像文件的格式
check_disk_format() {
    local disk_path="$1"
    if ! qemu-img info "$disk_path" >/dev/null 2>&1; then
        echo "错误：无法识别的磁盘格式。"
        exit 1
    fi
}

# 挂载虚拟磁盘到NBD设备
setup_nbd() {
    local disk_path="$1"
    echo "加载 nbd 模块..."

    # 尝试加载nbd模块
    if ! sudo modprobe nbd max_part=16 2>/dev/null; then
        echo "nbd模块加载失败，尝试安装内核扩展模块..."
        sudo apt-get install -y "linux-modules-extra-$(uname -r)"
        if ! sudo modprobe nbd max_part=16; then
            echo "错误：无法加载nbd模块！"
            echo "请手动执行：sudo apt install linux-modules-extra-$(uname -r)"
            exit 1
        fi
    fi

    NBD_DEVICE=$(select_free_nbd)
    if [ -z "$NBD_DEVICE" ]; then
        echo "未找到可用的 nbd 设备，退出。"
        exit 1
    fi

    echo "连接 $disk_path 到 $NBD_DEVICE..."
    sudo qemu-nbd -c "$NBD_DEVICE" "$disk_path"
    sleep 2  # 等待设备初始化
    VIRTUAL_DISK_ATTACHED=true
}

# 挂载所有分区
mount_all_partitions() {
    local device="$1"
    echo "检测并挂载所有分区..."

    # 强制重新读取分区表
    sudo partprobe "$device" >/dev/null 2>&1
    sleep 1

    # 获取所有分区
    local partitions=$(lsblk -rno NAME,TYPE "$device" | awk '$2 == "part" {print $1}')

    if [[ -z "$partitions" ]]; then
        echo "未找到任何分区！"
        cleanup
        exit 1
    fi

    for part in $partitions; do
        local partition="/dev/$part"
        local mount_point="$VDISK_MOUNT_POINT/$part"

        # 跳过已挂载的分区
        if findmnt -n -o TARGET "$partition" >/dev/null; then
            echo "分区 $partition 已挂载，跳过..."
            continue
        fi

        echo "挂载分区 $partition 到 $mount_point..."
        sudo mkdir -p "$mount_point"
        if ! sudo mount "$partition" "$mount_point"; then
            echo "挂载失败，跳过该分区。"
            continue
        fi
        MOUNT_POINTS+=("$mount_point")
    done

    if [[ ${#MOUNT_POINTS[@]} -eq 0 ]]; then
        echo "错误：没有成功挂载任何分区。"
        cleanup
        exit 1
    fi

    echo "已挂载的分区："
    mount | grep "$VDISK_MOUNT_POINT"
}

# 检查是否是根文件系统
is_root_partition() {
    local mount_point="$1"
    [[ -d "$mount_point/bin" && -d "$mount_point/etc" && -d "$mount_point/sbin" ]]
}

# 进入系统
enter_system() {
    local root_partition=""
    for mount_point in "${MOUNT_POINTS[@]}"; do
        if is_root_partition "$mount_point"; then
            root_partition="$mount_point"
            break
        fi
    done

    if [[ -n "$root_partition" ]]; then
        echo "检测到根分区: $root_partition，切换 chroot 环境..."
        chroot_to_root "$root_partition"
    else
        echo "未找到根文件系统，进入 bash 交互环境..."
        fallback_bash
    fi
}

# 进入chroot环境
chroot_to_root() {
    local root="$1"
    echo "Chrooting to $root..."
    sudo mount --bind /proc "$root/proc"
    sudo mount --bind /sys "$root/sys"
    sudo mount --bind /dev "$root/dev"
    sudo mount --bind /dev/pts "$root/dev/pts"
    sudo chroot "$root" /bin/bash
    sudo umount "$root/proc"
    sudo umount "$root/sys"
    sudo umount "$root/dev/pts"
    sudo umount "$root/dev"
}

# 进入备用shell
fallback_bash() {
    echo "在 $VDISK_MOUNT_POINT 下进入 bash shell 以手动管理挂载的分区..."
    pushd "$VDISK_MOUNT_POINT" > /dev/null
    bash --rcfile <(cat ~/.bashrc; echo 'PS1="(disk-mount) $PS1"')
    popd > /dev/null
}

# 获取用户输入
get_disk_input() {
    if [ -n "$1" ]; then
        DISK="$1"
    else
        echo "没有提供设备或镜像文件，现列出所有磁盘："
        lsblk -o NAME,MAJ:MIN,RM,LABEL,SIZE,FSTYPE,UUID,RO,TYPE,MOUNTPOINT,VENDOR,HOTPLUG -e7
        read -rp "请输入设备路径（如 /dev/sda）或虚拟磁盘镜像文件路径: " DISK
    fi
}

# 处理设备类型
process_disk() {
    if [ -b "$DISK" ]; then
        echo "选择的物理磁盘: $DISK"
        TARGET_DEVICE="$DISK"
    elif [ -f "$DISK" ]; then
        echo "选择的虚拟磁盘镜像文件: $DISK"
        IS_VIRTUAL=true
        check_disk_format "$DISK"
        setup_nbd "$DISK"
        TARGET_DEVICE="$NBD_DEVICE"
    else
        echo "错误： '$DISK' 既不是块设备也不是普通文件。"
        exit 1
    fi
}

# 主函数
main() {
    if [[ $# -ne 1 ]]; then
        echo "使用方法: $0 <磁盘镜像文件或块设备>"
        echo "示例:"
        echo "  $0 ./disk.img (虚拟磁盘)"
        echo "  $0 /dev/sdb  (物理磁盘)"
        exit 1
    fi

    install_dependencies
    get_disk_input "$1"
    process_disk
    mount_all_partitions "$TARGET_DEVICE"
    enter_system
}

main "$@"
